#include <windows.h>
#include <windef.h>
#include <winbase.h>
#include <winsvc.h>
#include <winnetwk.h>
#include <npapi.h>
#include <devioctl.h>

#include <strsafe.h>
#include <cassert>

#include "../hwfs/hwfs_shared.h"
#include "hwfs_np.h"
#include "debug.h"

#include "hwfs_enum.h"
#include <vector>

static inline int StringSizeBytes(const std::wstring& str)
{
    if (str.size() == 0) {
        return 0;
    }
    return (int) (str.size() + 1) * sizeof(WCHAR);
}

WCHAR* CopyStringToResInfo(WCHAR **resinfo_ptr, WCHAR* dst, const std::wstring& str,
    const UCHAR *endptr)
{
    if (str.size() == 0) {
        *resinfo_ptr = NULL;
        return dst;
    }

    *resinfo_ptr = dst;
    size_t maxSize = endptr - (UCHAR*)dst;
    int err = memcpy_s(dst, maxSize, str.c_str(), sizeof(WCHAR) * (str.size() + 1));
    assert(err == 0);
    return dst + str.size() + 1;
}


ULONG NPEnumEntry::FillResourceInfo(LPNETRESOURCE pResInfo, UCHAR** ppStringsEndPtr, ULONG maxSize)
{
    // compute all strings size
    ULONG stringsSize = StringSizeBytes(mLocalName);
    stringsSize += StringSizeBytes(mRemoteName);
    stringsSize += StringSizeBytes(mComment);
    stringsSize += StringSizeBytes(mProvider);

    const ULONG retSize = sizeof(*pResInfo) + stringsSize;
    if (maxSize < retSize) {
        return retSize;
    }

    pResInfo->dwScope = mScope;
    pResInfo->dwType = mType;
    pResInfo->dwDisplayType = mDisplayType;
    pResInfo->dwUsage = mUsage;

    // Move ppStringsEndPtr to the left of the buffer
    UCHAR* pEndPtr = *ppStringsEndPtr;
    *ppStringsEndPtr = pEndPtr - stringsSize;

    WCHAR* pw = (WCHAR *)*ppStringsEndPtr;
    pw = CopyStringToResInfo(&pResInfo->lpLocalName, pw, mLocalName, pEndPtr);
    pw = CopyStringToResInfo(&pResInfo->lpRemoteName, pw, mRemoteName, pEndPtr);
    pw = CopyStringToResInfo(&pResInfo->lpComment, pw, mComment, pEndPtr);
    pw = CopyStringToResInfo(&pResInfo->lpProvider, pw, mProvider, pEndPtr);

    DBG_PRINT((L"Resource: '%s' -- '%s' -- '%s'\n",
        mLocalName.c_str(), mRemoteName.c_str(), mComment.c_str()));

    return retSize;
}

NPEnum::NPEnum(DWORD dwScope, DWORD dwType, DWORD dwUsage, LPNETRESOURCE  lpNetResource)
    : mScope(dwScope),
    mType(dwType),
    mUsage(dwUsage)
{
    mIsRootResource = false;
    if (!lpNetResource || !lpNetResource->lpRemoteName || '\0' == lpNetResource->lpRemoteName[0]) {
        mIsRootResource = true;
    }

    std::wstring remoteName = L"<null>";
    if (lpNetResource && lpNetResource->lpRemoteName) {
        if ('\0' == lpNetResource->lpRemoteName[0]) {
            remoteName = L"<empty>";
        } else {
            remoteName = lpNetResource->lpRemoteName;
        }
    }

    // Zero matches all flags
    if (mUsage == 0)
        mUsage = RESOURCEUSAGE_CONNECTABLE | RESOURCEUSAGE_CONTAINER;

    DBG_PRINT((L"Enum Resources: scope:%08x, type:%08x, usage: %08x, name:%s\n",
        dwScope, dwType, dwUsage, remoteName.c_str()));
}

DWORD NPEnum::EnumerateResources()
{
    DWORD ret;

    mResourceList.clear();
    switch (mScope) {
        case RESOURCE_CONTEXT:
            ret = createHostEntry();
            break;
        case RESOURCE_CONNECTED:
            ret = createConnectedList();
            break;
        case RESOURCE_GLOBALNET:
            ret = createGlobalNetList();
            break;
        default:
            ret = WN_NOT_SUPPORTED;
            break;
    }

    return ret;
}

DWORD NPEnum::createHostEntry()
{
    NPEnumEntry entry(
        0, RESOURCETYPE_ANY, RESOURCEDISPLAYTYPE_SERVER, RESOURCEUSAGE_CONTAINER,
        NULL, HWFS_ROOTSRV_NAME, HWFS_ROOTSRV_COMMENT, HWFS_PROVIDER_NAME_U);

    mResourceList.push_back(entry);
    mResourcesListIterator = mResourceList.begin();

    return WN_SUCCESS;
}

const unsigned CONNECT_BUFFER_INIT_SIZE = 1024;

static DWORD HwfsGetConnectedList(std::vector<UCHAR>& ConnectionsBuffer)
{
    DWORD status = WN_MORE_DATA;
    while (status == WN_MORE_DATA) {
        ULONG copyBytes = (ULONG)ConnectionsBuffer.size();
        status = SendToMiniRdr(IOCTL_HWFS_GETLIST, NULL, 0, ConnectionsBuffer.data(), &copyBytes);
        DBG_PRINT((L"HWFS_GETLIST status 0x%x, copyBytes 0x%x\n", status, copyBytes));
        if (status == WN_MORE_DATA) {
            if (copyBytes <= ConnectionsBuffer.size()) {
                return WN_DEVICE_ERROR;
            } else {
                ConnectionsBuffer.resize(copyBytes);
            }
        } else if (status == 0) {
            ConnectionsBuffer.resize(copyBytes);
        }
    }
    return status;
}

DWORD NPEnum::createConnectedList()
{
    std::vector<UCHAR> connectionsBuffer(CONNECT_BUFFER_INIT_SIZE);
    DWORD status = HwfsGetConnectedList(connectionsBuffer);
    if (status) {
        DBG_PRINT((L"HWFS_GETLIST failed\n"));
        return status;
    }

    // We have multistring in ConnectionsBuffer.
    // each string has format \\device\miniredirector\;<DriveLetter>:\Server\Share
    WCHAR *wp = (WCHAR *)connectionsBuffer.data();
    WCHAR* end = wp + connectionsBuffer.size() / sizeof(WCHAR);
    DBG_PRINT((L"HWFS_GETLIST: %p-%p, *wp = 0x%x\n", wp, end, (unsigned)*wp));
    while (wp < end && '\0' != *wp) {
        NPEnumEntry entry;

        WCHAR* str = wp;
        size_t len = wcsnlen(str, end - wp);
        if ((wp + len) == end) {
            DBG_PRINT((L"HWFS_GETLIST returned invalid buffer\n"));
            return WN_DEVICE_ERROR;
        }

        std::wstring conn(str);
        size_t pos0 = conn.find(';');
        if (pos0 == std::wstring::npos) {
            DBG_PRINT((L"HWFS_GETLIST: invalid entry %s\n", str));
            return WN_DEVICE_ERROR;
        }
        pos0++; // pos0 points to the localname

        size_t pos1 = conn.find('\\', pos0);
        if (pos1 == std::wstring::npos) {
            DBG_PRINT((L"HWFS_GETLIST: invalid entry %s\n", str));
            return WN_DEVICE_ERROR;
        }

        entry.mLocalName = conn.substr(pos0, pos1 - pos0); // <DriveLetter>:
        entry.mRemoteName = std::wstring(L"\\\\") + conn.substr(pos1 + 1); // \\Server\Share
        entry.mProvider = HWFS_PROVIDER_NAME_U;
        entry.mDisplayType = RESOURCEDISPLAYTYPE_SHARE;
        entry.mScope = RESOURCE_CONNECTED;
        entry.mType = RESOURCETYPE_DISK;
        entry.mUsage = 0;
        mResourceList.push_back(entry);

        DBG_PRINT((L"push_entry: %s:%s\n", entry.mLocalName.c_str(), entry.mRemoteName.c_str()));

        wp += len + 1;
    }

    mResourcesListIterator = mResourceList.begin();

    return WN_SUCCESS;
}

DWORD NPEnum::createGlobalNetList()
{
    if (mIsRootResource) {
        return createHostEntry();
    }
    return createSharedFoldersList();
}

DWORD NPEnum::createSharedFoldersList()
{
    // We could fill all connected on host Folders.
    // But we return only connected ones.
    return createConnectedList();
}

DWORD NPEnum::FillResourceRequest(ULONG* pCount, VOID* pBuf, ULONG* pBufSize)
{
    if (mResourcesListIterator == mResourceList.end())
        return WN_NO_MORE_ENTRIES;

    DWORD BufSize = *pBufSize;
    // We have to fill strings starting from the end of the buffer,
    // because pBuf starts with array of fixed-size NETRESOURCE entries
    UCHAR* pStringsEndPtr = (UCHAR*)pBuf + BufSize;

    LPNETRESOURCE pResource = (LPNETRESOURCE)pBuf;
    for (*pCount = 0;
            mResourcesListIterator != mResourceList.end();
            mResourcesListIterator++, (*pCount)++, pResource++) {
        ULONG sz = mResourcesListIterator->FillResourceInfo(pResource, &pStringsEndPtr, BufSize);
        if (sz > BufSize) {
            // Return WN_MORE_DATA only if first entry doesn't fit into the buffer
            if (*pCount == 0) {
                *pBufSize = sz;
                return WN_MORE_DATA;
            }
            return WN_SUCCESS;
        }
        BufSize -= sz;
    }

    return WN_SUCCESS;
}
